package info.wangyuan.agent.service.event.listener;


import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import info.wangyuan.agent.common.Enum.FileEventType;
import info.wangyuan.agent.config.ModelFactory;
import info.wangyuan.agent.entity.dto.RagDocSyncChunkMsgDTO;
import info.wangyuan.agent.entity.dto.RagDocSyncStorageMsgDTO;
import info.wangyuan.agent.entity.po.DocumentUnit;
import info.wangyuan.agent.entity.po.FileInfo;
import info.wangyuan.agent.exception.BusinessException;
import info.wangyuan.agent.mapper.DocumentUnitMapper;
import info.wangyuan.agent.mapper.FileInfoMapper;
import info.wangyuan.agent.service.curd.CommonService;
import info.wangyuan.agent.service.event.RagFileEvent;
import info.wangyuan.agent.service.rag.RagDocSyncChunkContext;
import info.wangyuan.agent.service.rag.RagDocSyncChunkStrategy;
import info.wangyuan.agent.service.rag.RagEmbeddingDomainService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;
import org.springframework.transaction.event.TransactionPhase;
import org.springframework.transaction.event.TransactionalEventListener;

import java.util.ArrayList;
import java.util.List;

/**
 * 文件切分、向量化增、删、重建处理
 * 因为泛型擦除会导致监听不到事件，所以传参改成了RagFileEvent<?>，没写具体类型
 *
 * @author Albert
 * @since 2025-08-31 22:15:37
 */
@Slf4j
@Component
public class RagDocSyncListener {

    private final RagEmbeddingDomainService ragEmbeddingDomainService;
    private final ModelFactory modelFactory;
    private final FileInfoMapper fileInfoMapper;
    private final RagDocSyncChunkContext ragDocSyncChunkContext;
    private final DocumentUnitMapper documentUnitMapper;
    private final CommonService commonService;

    public RagDocSyncListener(RagEmbeddingDomainService ragEmbeddingDomainService, ModelFactory modelFactory, FileInfoMapper fileInfoMapper, RagDocSyncChunkContext ragDocSyncChunkContext, DocumentUnitMapper documentUnitMapper, CommonService commonService) {
        this.ragEmbeddingDomainService = ragEmbeddingDomainService;
        this.modelFactory = modelFactory;
        this.fileInfoMapper = fileInfoMapper;
        this.ragDocSyncChunkContext = ragDocSyncChunkContext;
        this.documentUnitMapper = documentUnitMapper;
        this.commonService = commonService;
    }

    /**
     * 文件切分事件
     */
    @Async
    @TransactionalEventListener(phase = TransactionPhase.AFTER_COMMIT)
    public void onFileChunking(RagFileEvent<?> evt) {
        if (evt.getType() != FileEventType.CHUNK) return;

        @SuppressWarnings("unchecked")
        var fileIds = (List<Integer>) evt.getPayload();
        log.info("开始文件切分处理，文件数量：{}", fileIds.size());
        fileInfoMapper.selectByIds(fileIds)
                .parallelStream() // 并行处理文件
                .forEach(this::processFileChunking);
    }

    /**
     * 文件批量向量化入库事件
     */
    @Async
    @TransactionalEventListener(phase = TransactionPhase.AFTER_COMMIT)
    public void onFileEmbeddingStorage(RagFileEvent<?> evt) {
        if (evt.getType() != FileEventType.STORAGE) return;

        @SuppressWarnings("unchecked")
        var payload = (List<RagDocSyncStorageMsgDTO>) evt.getPayload();
        ragEmbeddingDomainService.syncEmbeddingStorageByBatch(payload);
    }

    /**
     * 根据文件ID批量删除文件数据事件
     */
    @Async
    @TransactionalEventListener(phase = TransactionPhase.AFTER_COMMIT)
    public void onFileDataRemoveByFileId(RagFileEvent<?> evt) {
        if (evt.getType() != FileEventType.REMOVE_BY_FILE_ID) return;

        @SuppressWarnings("unchecked")
        var payload = (List<Integer>) evt.getPayload();
        // 删除文件单元
        commonService.deleteDocUnitsByFileIds(payload);
        // 删除文件向量
        ragEmbeddingDomainService.removeEmbeddingByFileIds(payload);
    }

    /**
     * 文件数据批量重建事件
     */
    @Async
    @TransactionalEventListener(phase = TransactionPhase.AFTER_COMMIT)
    public void onFileDataRebuildByFileId(RagFileEvent<?> evt) {
        if (evt.getType() != FileEventType.REBUILD_BY_FILE_ID) return;

        try {
            log.info("文件向量批量重建开始");
            @SuppressWarnings("unchecked")
            var payload = (List<Integer>) evt.getPayload();
            // 删除文件单元
            commonService.deleteDocUnitsByFileIds(payload);
            // 删除文件向量
            ragEmbeddingDomainService.removeEmbeddingByFileIds(payload);
            // 重新生成文件单元和向量
            onFileChunking(new RagFileEvent<>(FileEventType.CHUNK, payload));
            log.info("文件向量批量重建完成");
        } catch (Exception e) {
            log.error("文件向量批量重建失败");
        }
    }

    /**
     * 根据知识库ID批量删除文件数据事件
     */
    @Async
    @TransactionalEventListener(phase = TransactionPhase.AFTER_COMMIT)
    public void onFileDataRemoveByDatasetId(RagFileEvent<?> evt) {
        if (evt.getType() != FileEventType.REMOVE_BY_DATASET_ID) return;

        var payload = (Integer) evt.getPayload();

        List<Integer> fileIds = fileInfoMapper.selectList(
                        Wrappers.lambdaQuery(FileInfo.class)
                                .in(FileInfo::getKnowledgeBaseId, payload)
                ).stream()
                .map(FileInfo::getId)
                .toList();

        // 删除文件信息
        commonService.deleteFileInfosByFileIds(fileIds);
        // 删除文件单元
        commonService.deleteDocUnitsByFileIds(fileIds);
        // 删除文件向量
        ragEmbeddingDomainService.removeEmbeddingByFileIds(fileIds);
    }

    /**
     * 文件数据批量重建事件
     */
    @Async
    @TransactionalEventListener(phase = TransactionPhase.AFTER_COMMIT)
    public void onFileDataRebuildByDatasetId(RagFileEvent<?> evt) {
        if (evt.getType() != FileEventType.REBUILD_BY_DATASET_ID) return;

        try {
            log.info("文件向量批量重建开始");
            var payload = (Integer) evt.getPayload();

            List<Integer> fileIds = fileInfoMapper.selectList(
                            Wrappers.lambdaQuery(FileInfo.class)
                                    .in(FileInfo::getKnowledgeBaseId, payload)
                    ).stream()
                    .map(FileInfo::getId)
                    .toList();

            // 删除文件信息
            commonService.deleteFileInfosByFileIds(fileIds);
            // 删除文件单元
            commonService.deleteDocUnitsByFileIds(fileIds);
            // 删除文件向量
            ragEmbeddingDomainService.removeEmbeddingByFileIds(fileIds);
            // 重新生成文件单元和向量
            onFileChunking(new RagFileEvent<>(FileEventType.CHUNK, fileIds));
            log.info("文件向量批量重建完成");
        } catch (Exception e) {
            log.error("文件向量批量重建失败");
        }
    }

    /**
     * 文件切分处理
     */
    private void processFileChunking(FileInfo fileInfo) {
        Integer fileId = fileInfo.getId();
        try {
            log.info("正在处理文件切分，文件ID：{}", fileId);
            String fileSuffix = fileInfo.getSuffix();

            if (fileSuffix == null) throw new BusinessException("文件扩展名不能为空");

            RagDocSyncChunkStrategy strategy = ragDocSyncChunkContext.getTaskExportStrategy(fileSuffix.toUpperCase());
            if (strategy == null) throw new BusinessException("不支持的文件类型: " + fileSuffix);

            RagDocSyncChunkMsgDTO chunkMsg = RagDocSyncChunkMsgDTO.builder()
                    .fileId(fileId)
                    .pageSize(fileInfo.getPageSize())
                    .chatModel(modelFactory.getCustomOcrModel())
                    .build();
            strategy.handle(chunkMsg, fileSuffix);

            if (fileInfoMapper.selectById(fileId).getPageSize() > 0) {
                log.info("文件切分完成，文件ID：{}", fileId);
            }

            autoStartVectorization(fileInfo);
        } catch (Exception e) {
            log.error("文件切分失败，文件ID：{}", fileId, e);
        }
    }

    /**
     * 自动启动向量化处理
     *
     * @param fileInfo 文件实体
     */
    private void autoStartVectorization(FileInfo fileInfo) {
        Integer fileId = fileInfo.getId();
        try {
            log.info("自动启动文件向量化处理，文件ID：{}", fileId);

            // 检查是否有可用的文档单元进行向量化
            LambdaQueryWrapper<DocumentUnit> wrapper = Wrappers.lambdaQuery(DocumentUnit.class)
                    .eq(DocumentUnit::getFileId, fileId)
                    .eq(DocumentUnit::getIsChunked, true)
                    .eq(DocumentUnit::getIsVectored, false);
            List<DocumentUnit> documentUnits = documentUnitMapper.selectList(wrapper);

            if (documentUnits.isEmpty()) {
                log.warn("该文件尚无可用于向量化的片段: {}", fileId);
                return;
            }

            var customEmbeddingModel = modelFactory.getCustomEmbeddingModel();

            List<RagDocSyncStorageMsgDTO> ragDocSyncStorageMsgDTOList = new ArrayList<>();

            for (DocumentUnit documentUnit : documentUnits) {
                RagDocSyncStorageMsgDTO ragDocSyncStorageMsgDTO = RagDocSyncStorageMsgDTO.builder()
                        .id(documentUnit.getId())
                        .fileId(documentUnit.getFileId())
                        .fileName(fileInfo.getName())
                        .pageIndex(documentUnit.getPageIndex())
                        .content(documentUnit.getContent())
                        .isVector(true)
                        .datasetId(fileInfo.getKnowledgeBaseId())
                        .embeddingModel(customEmbeddingModel)
                        .build();
                ragDocSyncStorageMsgDTOList.add(ragDocSyncStorageMsgDTO);
            }

            // 文件单元批量向量化入库
            onFileEmbeddingStorage(new RagFileEvent<>(FileEventType.STORAGE, ragDocSyncStorageMsgDTOList));

            log.info("为文件： {} 自动启动向量化处理, 共 {} 个文件片段 ", fileId, documentUnits.size());

        } catch (Exception e) {
            log.error("为文件： {} 自动启动向量化处理失败", fileId, e);
            // 如果自动启动失败，重置向量化状态
        }
    }
}
